function [outSt,logL]=kfilterNanReg(parvec,funcmod, Y, trainvec, solveopt, addsol)
% =========================================================================
% kfilterNanReg.m 
% Model with NANs, but no splits nor cummulators. 
%
% KFILTER_TAGIN_FLEX
% function [stt,etamat,smooth_st,yfor,vstar,logL,likelvec]=kfilter_tagin(parest, parvec, parposest, funcmod, y, trainvec, solveopt, addsol)
% 
% =========================================================================
%% I. Purpose
% 
% Kalman Filter with missing observations. Other than the mapping from
% observables to states, the matrices are assumed time invariant. 
% Obtains smoothed estimates of the states and disturbances 
% using the disturbance smoother in Durbin and Koopman 
%
% Model 
%
% $$ y_{t}   = Z*s_{t}   + Z*CC $$
%
% $$ s_{t}   = G*s_{t-1} + R*n_{t} $$ 
% 
% $$ V( n_{t} ) = SDX'*SDX $$
% 
% This does NOT require a deterministc pattern in the variation of Z, i.e.
% code will check for NANs at each t and consruct 
% 
% $$ Z_{t}=W_{t}*Z $$ 
% 
% by extracting the appropriate non-NAN rows 
%
% 
%% II. Inputs 
% 
% *PAREST*: vector of parameters 
% 
% *FUNCMOD*: model function handle 
% 
% *Y*: [nobs nz] data vector, with NANS 
% 
% *TRAINVEC:* [2 1] vector indicating which row to start in and finish for
% the evaluation of the likelihood. Note: if empty [1 nobs] will be used. 
% 
% *SOLVEOPT*: setting for solution, e.g. for NL solver 
% 
% *ADDSOL*: additional input needed to solve model 
% 
%% III. Outputs 
% % 
% *stt*: [nobs ns] vector of one-sided states 
%
% *etamat*: [nobs nx] vector of smooth disturbances 
%
% *smooth_st*: [nobs ns] vector of smooth states 
% 
% *yfor*: [nobs nz] vector of one step ahead forecasts for observables. 
% Note: notation is not exactly as in data unless NANs all grouped at the 
% end of the observation vector
% 
% *vfor*: [nobs nz] vector of one step ahead forecast errors for observables. 
% Note: same as yfor, notation is not exactly as in data unless NANs all grouped at the 
% end of the observation vector
%
% *logL*: [1 1] log likelihood with constant included 
% 
% *countmat*: [nobs ns nx] decomposition of all states when 1 shock at a
% time turned on. 1 shock per page. 
%
% *counobs*: [nobs nz nx] decomposition of all obserbvables when 1 shock at a
% time turned on. 1 shock per page. 
% 
% *countszero*: [ns nx] decomposition of state S(1) when 1 shock at a time
% turned on. 1 shock per column. 
% 
% *likelvec*: [nobs 1] vector of conditional likelihoods 
% 
%% IV. Notes 
%  
% Code will check that can recover the observables through a forward state
% smoothing recursion. Will not succeed if the SDX matrix is not diagonal.
% In this case a warning message will be issued. If the SDX matrix is
% diagonal, then error message instead. 
% 
% June 30th 2010: Checked Identical results to *KFILTER_TAGIN.m* for models with M
% to Q patter in Z variation.
% 
%% V. Related Codes 
% 
% *KFILTER_TAGIN.m* : Assumes M to Q pattern in variation of Z, requires
% additional outputs from model solution 
% 
% *KFILTER_TAGTV.m* : For G time varying deterministically with a
% cummulator 
%
%
% Alejandro Justiniano 
% Created: June 30th 2010  
% =========================================================================

%% 1. Solve model 
[G,R,C,eu,SDX,Z,structOne,ssvec,structTwo]=feval(funcmod,parvec,solveopt,addsol);
if ~isequal(eu,[1;1]);return;end 

%% 2. Initial Dimensions 
[T,ny]=size(Y);
[ns,nx]  =size(R);

if  isempty(trainvec); 
    disp('Default TRAINVEC: use all sample') 
    trainvec=[1 nobs]; 
end 

%% 3. Demean Data using model consistent variables 
if any(C~=0)
    disp('Model demeaning the data');
    Y=(Y-repmat((Z*C)',[T 1]))';
else
    Y=Y';
end

%% 4. Initial states and Initial dispersion 
azero=addsol.aZero; %zeros(ns,1);
pshat=addsol.pZero; %lyapunov_symm(G,R*(SDX')*SDX*R');
P0=pshat; 

pnx=zeros(ns,ns,nx);
for ii=1:nx
    pnx(:,:,ii)=lyapunov_symm(G,R(:,ii)*SDX(ii,ii)*SDX(ii,ii)*R(:,ii)');
end

%% 5.Storage Matrices 
vstar=nan(ny,T); 
finmat=nan(ny,ny,T); 
kpartg=nan(ns,ny,T); 
likelvec=zeros(T,1); 
yfor=nan(ny,T);
% Matrices with one additional entry (initialization)
stt=zeros(ns,T+1);  
ptt=zeros(ns,ns,T+1); 
stt(:,1)=azero;  
ptt(:,:,1)=pshat; 

Zdim      =zeros(T,1); 
mat_obspos=zeros(T,ny); 

%% 6.Forward Filter 

W   =eye(ny);
for ii=1:T;

    %% 6.A Handling of missing observations 
    ytt=Y(:,ii);
    
    % Determine W and position of the NAN  
    ind =~isnan(ytt);
    rowt =find(~isnan(ytt));
    
    %% 6.B ZDIM, MAT_OBSPOS, DIMT 
    % ZDIM: stores the number of observable series at each t 
    %
    % MAT_OBSPOS: stores position of the rows of Y that are observed 
    % 
    % DIMT: Vector with (1:ZDIM(ii)) indicator 
    %
    % In this way, the $$ Z_{t}=W_{t}Z $$ matrix can be easily recovered as
    % W(t)=W( MAT_OBSPOS(ii,1:ZDIM(ii)),: ) 
    % 
    ytt=ytt(ind);    
    Zdim(ii)=length(ytt); 
    Ztt=W((ind==1),:)*Z;
    mat_obspos(ii,1:Zdim(ii))=rowt; 
    dimt=( 1:Zdim(ii) ); 
    
    yfor(dimt,ii)=Ztt*G*stt(:,ii);
    [stt(:,ii+1),ptt(:,:,ii+1),likelvec(ii),vvv,fin,kgain]=...
        kf_dk(ytt,Ztt,stt(:,ii),ptt(:,:,ii),G,R*SDX');    
    % Output dimension depends on DIMT 
    vstar(dimt,ii)=vvv; 
    finmat(dimt,dimt,ii)=fin; 
    kpartg(:,dimt,ii)=kgain;     
    
end
logL=-0.5*sum(Zdim)*log(2*pi)+sum(likelvec(trainvec(1):trainvec(2)));

% Truncate First Observation which is the initialization 
stt  =stt(:,2:end); 
ptt  =ptt(:,:,2:end); 
yfor=yfor'; 

%% 7.Backward Filter 
% Obtain the Innovation using a disturbance smoother 
% Use same recursion to obtain STATE smoother 

%% 7.A Storage of the backward recurions
etamat=zeros(nx,T); 
smooth_st=zeros(ns,T);  
rstar=zeros(ns,1);
rmat=zeros(ns,T);
% First smooth state
[rstar,etamat(:,end)]=smoothdis(rstar,SDX'*SDX,(R'),(Ztt)',finmat(dimt,dimt,end),zeros(ns),vstar(dimt,end));
smooth_st(:,ii)=stt(:,end)+ptt(:,:,end)*rstar;
rmat(:,end)=rstar;
for ii=(T-1):-1:1;
    %% 7.B Varying dimension in the backward recurions
    % 1. DIMT
    dimt=( 1:Zdim(ii) );
    % 2. Ztt=Wtt*Z
    Ztt=( W( mat_obspos(ii,1:Zdim(ii)),:) )*Z;
    %Adjust the dimension of the Z matrix accordingly
    [rstar,etamat(:,ii)]=smoothdis(rstar,SDX'*SDX,(R'),...
        (Ztt'),finmat(dimt,dimt,ii),((G-G*kpartg(:,dimt,ii)*Ztt)'),vstar(dimt,ii));
    smooth_st(:,ii)=stt(:,ii)+ptt(:,:,ii)*rstar;
    rmat(:,ii)=rstar;
end
a0 = P0*(G')*rstar;   % Note: this is only correct in the case where a0 = zeros(ns,1) 
etamat=etamat';
smooth_st=smooth_st';
vstar=vstar';


% outSt.sttNoInitial=stt; 
% outSt.pttNoInitial=ptt; 
%% Initial state truncated above!
outSt.aZero=a0; 
outSt.aLast=stt(:,end);
outSt.pLast=ptt(:,:,end); 


% Check there are NO NANS in the smooth states and innovations 
if sum( any( isnan(smooth_st) ) ) > 0 
    error('There are NANs in smoothed states') 
end 
if sum( any( isnan(etamat) ) ) > 0 
    error('There are NANs in smoothed innovations') 
end 


%% 7. Check Smoother 
% Check that Smooth States are identical if using disturbance smoother (above) vs. state smoother (below) 
% and if can also recover the observables
tol=1e-5; 
% State at time zero give by GG1*azero+Pzero*rstar;
[yCheck,smoothCheck]=kfilterRegSplitSimulation(a0,etamat'); 
disp('Max Discrepancy Smooth vs. Actual Data'); 
maxdifY=comparemat(yCheck,Y'); 
disp('Max Discrepancy State and Innovation Smoother'); 
maxdifS=comparemat(smoothCheck,smooth_st); 
if maxdifY > tol || maxdifS > tol 
    error('Smoother discrepancy exceeds tolerance') 
end 

%% 8. Initial Variance and State per shock (used below for the counterfactual decompositions)
% vInitialPershock: Initial Variance, decomposed per shock 
% sOnePerShock, State at time 1 decomposed, decomposed per shock
vInitialPerShock=zeros(ns,ns,nx);
sZeroSmoothPerShock=zeros(ns,nx); 
sOneSmoothPerShock =zeros(ns,nx); 
for ii=1:nx
    vInitialPerShock(:,:,ii)=lyapunov_symm(G,R(:,ii)*SDX(ii,ii)*SDX(ii,ii)*(R(:,ii)'));
    sZeroSmoothPerShock(:,ii)=vInitialPerShock(:,:,ii)*G'*rstar;
    etatemp = zeros(nx,1);
    etatemp(ii) = etamat(1,ii);
    sOneSmoothPerShock(:,ii)=G*sZeroSmoothPerShock(:,ii)+R*etatemp;
end
maxdifSzero=comparemat(sum(sZeroSmoothPerShock,2),a0); 
maxdifSone =comparemat(sum(sOneSmoothPerShock,2),smooth_st(1,:)' ); 
if maxdifSzero > tol; error('Cannot decompose sZeroPerShock'); end 
if maxdifSone > tol; error('Cannot decompose sOnePerShock'); end


%% 9. Counterfactual Decomposition 
% Generate states by feeding each shock at a time, ensure that it recovers
% the original state 
disp('Begin Counterfactual')
countStates=zeros(T,ns,nx);
countObs   =zeros(T,ny,nx);
for ii=1:nx
    etaTemp=zeros(T,nx); 
    etaTemp(:,ii)=etamat(:,ii); 
    [countObs(:,:,ii),countStates(:,:,ii)]=kfilterRegSplitSimulation(sZeroSmoothPerShock(:,ii),etaTemp');
end 
disp('Maximum Discrepancy Counterfactual States and Smooth States') 
maxdifCount=comparemat(sum(countStates,3),smooth_st); 
if maxdifCount > tol; error('Counterfactual States do not recover smooth states'); end 
    
outSt.filteredSt=stt'; 
outSt.smoothSt=smooth_st; 
outSt.innovations=etamat; 
outSt.countSt=countStates; 
outSt.countObs=countObs; 
outSt.decompInitialSt=sOneSmoothPerShock; 
outSt.forecastObs=yfor; 

   
%% Subroutine kfilterRegSplitSimulation allows to simulate the model  
% Inputs 
% sInitial: [NS 1] Initial State 
% innovMat: [Nx Nobx] matrix of innovations 
% By being a sub-routine it has access to all variables defined above 
% Be-careful not to use an index (ii,jj) used above or to repeat variable
% names
    function [ySim,sSim]=kfilterRegSplitSimulation(sInitial,innovMat) 
        if ~isequal(size(innovMat),[nx T]) 
            error('Input innovMat must be [NX Nobs]') 
        end 
        sSim=zeros(ns,T);
        ySim=zeros(size(Y));
        sSim(:,1)=G*sInitial+R*innovMat(:,1);
        ySim(:,1)=Z*sSim(:,1);
        for kk=2:T;
            sSim(:,kk)=G*sSim(:,kk-1)+R*innovMat(:,kk);
            ySim(:,kk)=Z*sSim(:,kk);
        end      
        ySim=ySim'; 
        sSim=sSim'; 
    end

%% End of File
end